import matplotlib.pyplot as plt
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor,plot_tree

# 构建数据
print(np.array(list(range(1, 11))))
x = np.array(list(range(1,11))).reshape(-1,1)
print(x)
print(x.shape)
y = np.array([5.56,5.7,5.91,6.4,6.8,7.05,8.9,8.7,9,9.05])
print(y)

# 模型训练
model1 = LinearRegression()
model2 = DecisionTreeRegressor(max_depth=1)
model3 = DecisionTreeRegressor(max_depth=3)

model1.fit(x,y)
model2.fit(x,y)
model3.fit(x,y)

# 模型预测
x_test = np.arange(0.0,10.0,0.01).reshape(-1,1)
print(x_test.shape)
y1 = model1.predict(x_test)
y2 = model2.predict(x_test)
y3 = model3.predict(x_test)

# 可视化
plt.scatter(x,y)
plt.plot(x_test,y1)
plt.plot(x_test,y2)
plt.plot(x_test,y3)
plt.grid()
plt.show()

plt.figure(figsize=(30,20))
plot_tree(model3,filled=True)
plt.show()

